import torch
import torch.nn as nn


class WeightHidden(nn.Module):
    def __init__(self, hidden_size):
        super().__init__()
        self.hidden_size = hidden_size
        self.dense(hidden_size, 1)

    def forward(self, inputs):
        inputs = torch.stack([inputs], 0)
        weights = self.dense(inputs).softmax(0)
        return (inputs * weights).sum(0)

    @property
    def outdim(self):
        return self.hidden_size

    def __repr__(self):
        return 'weighted_hidden'
